import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from modules.yolov3.yololayer import YoloLayer
from torchvision.models import resnet
from torchvision.models import mobilenet_v2

# <function make_conv2d_bn_leaky/>
def make_conv2d_bn_leaky(in_channels, out_channels, kernel_size, stride=1, padding=0, dilation=1, groups=1, padding_mode='zeros'):
    return nn.Sequential(
        nn.Conv2d(in_channels, out_channels, kernel_size, stride, padding, dilation, groups, False, padding_mode),
        nn.BatchNorm2d(out_channels, eps=1e-5, momentum=0.1, affine=True, track_running_stats=True),
        nn.LeakyReLU()
    )
# </function make_conv2d_bn_leaky>

# <class FaceMobileNet(nn.Module)/>
class FaceMobileNet(nn.Module):
    """Some Information about FDNet. Data: 2019-09-23 11:30 """
    # <method __init__/>
    def __init__(
        self, 
        grids,
        anchors,
        num_classes,
        img_size,
        ):
        super(FaceMobileNet, self).__init__()
        self._feature = mobilenet_v2(pretrained=True).features
        for name, value in self._feature.named_parameters():
            value.requires_grad = False
            print("set parameter {0}'s requires_grad to Fasle".format(name))
        # end-for
        for mod in self._feature.modules():
            if mod.__class__.__name__ in ["ReLU6", "ReLU"]:
                mod = nn.LeakyReLU()
                print("replace the relu6 to leayrelu")
            # end-if            
        # end-for
        for mod in self._feature.modules():
            if mod.__class__.__name__ == nn.ReLU6.__class__.__name__:
                print("still had relu6")
            # end-if
        # end-for        
        
        # out layers ...
        self._out = nn.Sequential(
            make_conv2d_bn_leaky(in_channels=1280, out_channels=512, kernel_size=3, stride=1, padding=1, dilation=1, groups=1),
            make_conv2d_bn_leaky(in_channels=512,  out_channels=256, kernel_size=3, stride=1, padding=1, dilation=1, groups=1),
            make_conv2d_bn_leaky(in_channels=256,  out_channels=128, kernel_size=3, stride=1, padding=1, dilation=1, groups=1),
            nn.Conv2d(in_channels=128, out_channels=len(anchors[0])*(5+num_classes), kernel_size=1, stride=1, padding=0, dilation=1, groups=1)
            )
        
        # yolo layer ...
        self._yolo = YoloLayer(anchors = anchors[0], num_classes = num_classes, image_size = img_size)
                
        # initialize params ...
        for m in self._out.modules():
            if isinstance(m, nn.Conv2d):
                nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='leaky_relu')
            elif isinstance(m, (nn.BatchNorm2d, nn.GroupNorm)):
                nn.init.constant_(m.weight, 1)
                nn.init.constant_(m.bias, 0)
            # end-if
        # end-for
    # </method __init__>

    # <method forward/>
    def forward(self, x):
        nI_h = x.size(2)
        nI_w = x.size(3)
        feature = self._feature(x)
        out = self._out(feature)
        encoded_fm, decode_fm = self._yolo(out, nI_h, nI_w)
        return encoded_fm, decode_fm
    # </method forward>
# </class FaceMobileNet(nn.Module)>

if __name__ == "__main__":
    import sys 
    sys.path.append('../ModelCompression')
    from graph import reconstructor
    from utils.compressmethod import showFMLAs
    
    grids = [[13,13]]
    anchors = [[[10,13],[16,30],[33,23],[30,61],[62,45],[59,119],[116,90],[156,198],[373,326]]]
    num_classes = 2
    inp = torch.rand(1, 3, 416, 416)        
    net = FaceMobileNet(grids, anchors, num_classes, (416,416))
    outp = net(inp)
    
    import pdb; pdb.set_trace()
    fmlas_origin = showFMLAs(inp, net)
    print("""FMLAs: {} M""".format(fmlas_origin / 1e6))